from fastapi import FastAPI
from fastapi_mcp import FastApiMCP
from pydantic import BaseModel
from typing import List

app = FastAPI(title="多MCP服务器示例")

# ===================== 数据模型 =====================
class NumbersInput(BaseModel):
    numbers: List[float]

class TwoNumbersInput(BaseModel):
    a: float
    b: float

class TextInput(BaseModel):
    text: str

# ===================== 计算器 API =====================
@app.post("/calculator/add", operation_id="calc_add", tags=["calculator"])
async def add_numbers(input: NumbersInput):
    """加法运算"""
    return {"result": sum(input.numbers), "operation": "addition"}

@app.post("/calculator/subtract", operation_id="calc_subtract", tags=["calculator"])
async def subtract_numbers(input: TwoNumbersInput):
    """减法运算"""
    return {"result": input.a - input.b, "operation": "subtraction"}

@app.post("/calculator/multiply", operation_id="calc_multiply", tags=["calculator"])
async def multiply_numbers(input: NumbersInput):
    """乘法运算"""
    result = 1
    for num in input.numbers:
        result *= num
    return {"result": result, "operation": "multiplication"}

@app.post("/calculator/divide", operation_id="calc_divide", tags=["calculator"])
async def divide_numbers(input: TwoNumbersInput):
    """除法运算"""
    if input.b == 0:
        return {"error": "不能除以零"}
    return {"result": input.a / input.b, "operation": "division"}

# ===================== 文本处理 API =====================
@app.post("/text/uppercase", operation_id="text_upper", tags=["text"])
async def text_uppercase(input: TextInput):
    """转换为大写"""
    return {"result": input.text.upper(), "operation": "uppercase"}

@app.post("/text/lowercase", operation_id="text_lower", tags=["text"])
async def text_lowercase(input: TextInput):
    """转换为小写"""
    return {"result": input.text.lower(), "operation": "lowercase"}

@app.post("/text/reverse", operation_id="text_reverse", tags=["text"])
async def text_reverse(input: TextInput):
    """反转文本"""
    return {"result": input.text[::-1], "operation": "reverse"}

@app.post("/text/count", operation_id="text_count", tags=["text"])
async def text_count(input: TextInput):
    """统计字符数"""
    return {
        "length": len(input.text),
        "words": len(input.text.split()),
        "operation": "count"
    }

# ===================== 数学函数 API =====================
@app.post("/math/power", operation_id="math_power", tags=["math"])
async def math_power(input: TwoNumbersInput):
    """计算幂次方"""
    return {"result": input.a ** input.b, "operation": "power"}

@app.post("/math/sqrt", operation_id="math_sqrt", tags=["math"])
async def math_sqrt(input: dict):
    """计算平方根"""
    import math
    num = input.get("number", 0)
    if num < 0:
        return {"error": "不能对负数开平方"}
    return {"result": math.sqrt(num), "operation": "sqrt"}

# ===================== 创建多个 MCP 服务器 =====================

# MCP 服务器 1: 计算器服务（只包含基础计算功能）
mcp_calculator = FastApiMCP(
    app,
    name="Calculator MCP Server",
    description="提供基础的加减乘除计算功能",
    include_operations=["calc_add", "calc_subtract", "calc_multiply", "calc_divide"]
)

# MCP 服务器 2: 文本处理服务（只包含文本处理功能）
mcp_text = FastApiMCP(
    app,
    name="Text Processing MCP Server",
    description="提供文本处理功能：大小写转换、反转、统计等",
    include_operations=["text_upper", "text_lower", "text_reverse", "text_count"]
)

# MCP 服务器 3: 数学服务（只包含高级数学功能）
mcp_math = FastApiMCP(
    app,
    name="Math MCP Server",
    description="提供高级数学计算功能",
    include_operations=["math_power", "math_sqrt"]
)

# MCP 服务器 4: 全功能服务（包含所有功能）
mcp_all = FastApiMCP(
    app,
    name="All-in-One MCP Server",
    description="包含所有计算器、文本和数学功能",
    include_tags=["calculator", "text", "math"]  # 使用 tags 过滤
)

# ===================== 挂载到不同的端点 =====================
mcp_calculator.mount_http(mount_path="/mcp/calculator")  # http://localhost:8000/mcp/calculator
mcp_text.mount_http(mount_path="/mcp/text")              # http://localhost:8000/mcp/text
mcp_math.mount_http(mount_path="/mcp/math")              # http://localhost:8000/mcp/math
mcp_all.mount_http(mount_path="/mcp/all")                # http://localhost:8000/mcp/all

# 也可以同时支持 SSE 传输
# mcp_calculator.mount_sse(mount_path="/sse/calculator")
# mcp_text.mount_sse(mount_path="/sse/text")

# ===================== 启动信息 =====================
print("=" * 80)
print("🚀 多 MCP 服务器已启动")
print("=" * 80)
print("\n📊 MCP 服务器列表:")
print(f"\n1. 计算器服务 (Calculator)")
print(f"   - 端点: http://localhost:8000/mcp/calculator")
print(f"   - 工具数量: {len(mcp_calculator.tools)}")
print(f"   - 工具: {[t.name for t in mcp_calculator.tools]}")

print(f"\n2. 文本处理服务 (Text Processing)")
print(f"   - 端点: http://localhost:8000/mcp/text")
print(f"   - 工具数量: {len(mcp_text.tools)}")
print(f"   - 工具: {[t.name for t in mcp_text.tools]}")

print(f"\n3. 数学服务 (Math)")
print(f"   - 端点: http://localhost:8000/mcp/math")
print(f"   - 工具数量: {len(mcp_math.tools)}")
print(f"   - 工具: {[t.name for t in mcp_math.tools]}")

print(f"\n4. 全功能服务 (All-in-One)")
print(f"   - 端点: http://localhost:8000/mcp/all")
print(f"   - 工具数量: {len(mcp_all.tools)}")
print(f"   - 工具: {[t.name for t in mcp_all.tools]}")

print("\n" + "=" * 80)
print("📖 API 文档: http://localhost:8000/docs")
print("=" * 80)

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)